import requests
import os
from dotenv import load_dotenv
# 加载环境变量
load_dotenv()

def check_api_permissions(api_key):
    headers = {"Authorization": f"Bearer {api_key}"}

    # 测试Chat权限
    chat_url = "https://open.bigmodel.cn/api/paas/v4/chat/completions"
    chat_resp = requests.post(
        chat_url,
        headers=headers,
        json={"model": "glm-4", "messages": [{"role": "user", "content": "test"}]},
        timeout=5
    )

    # 测试Embedding权限
    embed_url = "https://open.bigmodel.cn/api/paas/v4/embeddings"
    embed_resp = requests.post(
        embed_url,
        headers=headers,
        json={"model": "text_embedding", "input": ["test"]},
        timeout=5
    )

    return {
        "chat": chat_resp.status_code == 200,
        "embedding": embed_resp.status_code == 200
    }

# 使用方式
permissions = check_api_permissions(os.getenv("ZHIPUAI_API_KEY"))
print(f"权限状态: 文本生成={permissions['chat']}, 向量嵌入={permissions['embedding']}")